import os
import sys
import time
import json
import redis
import random
import threading
from paho.mqtt import client as mqtt_client
from utils.custom_logger import CustomLogger
from settings import (MQTT_HOST, MQTT_PORT, MQTT_CLIENT_CERT, MQTT_AUTH, MQTT_PASSWORD, MQTT_TOPIC, 
                      MQTT_RECEIVE_QUEUE, MQTT_SEND_QUEUE,
                      REDIS_HOST, REDIS_PORT, REDIS_DB
                      )

class MqttSubscriber:
    logger = CustomLogger(
        logger_name='mqtt_module', 
        log_file_path='./logs/mqtt_subscriber.log'
    ).get_logger()

    def __init__(self):
        self.broker = MQTT_HOST
        self.port = MQTT_PORT
        self.username = MQTT_AUTH
        self.password = MQTT_PASSWORD
        self.topic = MQTT_TOPIC
        self.client_cert = MQTT_CLIENT_CERT
        self.client_id = f'python-mqtt-{random.randint(0, 100)}'
        self.client = self._create_client()
        self._ensure_logs_directory_exists()

        self.redis_client = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB)

    def _ensure_logs_directory_exists(self):
        log_dir = os.path.dirname('./logs/mqtt_module.log')
        if not os.path.exists(log_dir):
            os.makedirs(log_dir)
            self.logger.info(f"Created logs directory: {log_dir}")

    def _create_client(self):
        client = mqtt_client.Client(self.client_id)
        try:
            client.tls_set(ca_certs=self.client_cert)
            client.tls_insecure_set(False)  # 启用严格证书验证
        except Exception as e:
            self.logger.error(f"TLS configuration failed: {e}")
            raise
        
        client.username_pw_set(self.username, self.password)
        client.on_connect = self._on_connect
        client.on_message = self._on_message
        client.on_disconnect = self._on_disconnect
        return client

    def _on_connect(self, client, userdata, flags, rc):
        if rc == 0:
            self.logger.info(f"Connected to {self.broker}:{self.port}")
            client.subscribe(self.topic)
            self.logger.info(f"Subscribed to topic: {self.topic}")
        else:
            error_map = {
                1: "Incorrect protocol version",
                2: "Invalid client identifier",
                3: "Server unavailable",
                4: "Bad username or password",
                5: "Not authorised"
            }
            self.logger.error(f"Connection failed ({rc}): {error_map.get(rc, 'Unknown error')}")
            # 连接失败时退出程序
            import sys
            sys.exit(1)

    def _on_message(self, client, userdata, msg):
        try:
            payload = msg.payload.decode()
            self.logger.info(f"Received [{msg.topic}]: {payload}")
            # 在这里可以添加消息处理逻辑 将数据存入redis队列
            self.redis_client.lpush(MQTT_RECEIVE_QUEUE, payload)
        except UnicodeDecodeError:
            self.logger.warning(f"Non-UTF8 payload received from {msg.topic}")

    def _on_disconnect(self, client, userdata, rc):
        if rc != 0:
            self.logger.warning(f"Unexpected disconnection (rc={rc})")

    def connect(self):
        try:
            self.logger.info(f"Connecting to MQTT broker: {self.broker}:{self.port}")
            self.client.connect(self.broker, self.port)
        except Exception as e:
            self.logger.critical(f"Connection error: {e}")
            # 致命错误时退出程序
            import sys
            sys.exit(1)

    def start(self):
        self.connect()
        try:
            self.logger.info("Starting MQTT client loop...")
            self.client.loop_forever()
        except KeyboardInterrupt:
            self.logger.info("Keyboard interrupt received, stopping...")
            self.stop()
            import sys
            sys.exit(0)  # 正常退出
        except Exception as e:
            self.logger.critical(f"Unexpected error in client loop: {e}")
            self.stop()
            import sys
            sys.exit(1)  # 异常退出

    def stop(self):
        self.logger.info("Stopping MQTT client...")
        self.client.disconnect()
        self.client.loop_stop()
        self.logger.info("MQTT client stopped successfully")



class MqttPublisher:
    logger = CustomLogger(
        logger_name='mqtt_publisher',
        log_file_path='./logs/mqtt_publisher.log'
    ).get_logger()

    def __init__(self):
        self.broker = MQTT_HOST
        self.port = MQTT_PORT
        self.username = MQTT_AUTH
        self.password = MQTT_PASSWORD
        self.client_cert = MQTT_CLIENT_CERT
        self.client_id = f'python-mqtt-publisher-{random.randint(0, 100)}'
        self.client = self._create_client()
        self._ensure_logs_directory_exists()
        
        # Redis 连接（添加异常处理）
        try:
            self.redis_client = redis.Redis(
                host=REDIS_HOST,
                port=REDIS_PORT,
                db=REDIS_DB,
                decode_responses=True  # 自动解码为字符串
            )
            self.redis_client.ping()  # 测试连接
        except redis.RedisError as e:
            self.logger.critical(f"Redis connection failed: {e}")
            raise
        
        self.running = False  # 运行状态标识
        self.queue_thread = None  # Redis 队列监听线程

    def _ensure_logs_directory_exists(self):
        log_dir = os.path.dirname('./logs/mqtt_publisher.log')
        if not os.path.exists(log_dir):
            os.makedirs(log_dir, exist_ok=True)  # 添加 exist_ok 避免并发创建冲突
            self.logger.info(f"Created logs directory: {log_dir}")

    def _create_client(self):
        client = mqtt_client.Client(self.client_id)
        try:
            client.tls_set(ca_certs=self.client_cert)
            client.tls_insecure_set(False)
        except Exception as e:
            self.logger.error(f"TLS configuration failed: {e}")
            raise
        
        client.username_pw_set(self.username, self.password)
        client.on_connect = self._on_connect
        client.on_disconnect = self._on_disconnect  # 移除无用的 on_message
        return client

    def _on_connect(self, client, userdata, flags, rc):
        if rc == 0:
            self.logger.info(f"Connected to MQTT Broker: {self.broker}:{self.port}")
        else:
            error_map = {
                1: "Incorrect protocol version",
                2: "Invalid client identifier",
                3: "Server unavailable",
                4: "Bad username or password",
                5: "Not authorised"
            }
            self.logger.error(f"Connection failed ({rc}): {error_map.get(rc, 'Unknown error')}")
            # 可选：自动重试连接（非阻塞方式）
            # self.client.reconnect()
            raise ConnectionError("MQTT connection failed")

    def _on_disconnect(self, client, userdata, rc):
        if rc != 0:
            self.logger.warning(f"Unexpected disconnection (rc={rc}). Reconnecting...")
            self.client.reconnect()  # 自动重连

    def start(self):
        """启动服务：连接 MQTT + 启动 Redis 队列监听线程"""
        try:
            self.client.connect(self.broker, self.port)
            self.client.loop_start()  # 启动非阻塞网络循环

            self.running = True
            self.queue_thread = threading.Thread(
                target=self._listen_redis_queue,
                name="redis-queue-listener"
            )
            
            self.queue_thread.start()

            self.logger.info(f"MQTT Publisher started. Listening to Redis queue: {MQTT_SEND_QUEUE}")

            self.queue_thread.join()  # 等待队列监听线程结束（主线程退出时自动结束）


        except KeyboardInterrupt as e:
            self.logger.info("Keyboard interrupt received, stopping...")
            self.stop()
            sys.exit(0)  # 正常退出

        except Exception as e:
            self.logger.critical(f"Service startup failed: {e}")
            self.stop()

    def _listen_redis_queue(self):
        """独立线程：监听 Redis 队列并发布消息"""
        while self.running:
            try:
                # blpop 返回 (queue_name, message) 元组，
                queue_name, task_data = self.redis_client.blpop(MQTT_SEND_QUEUE)
                
                if task_data:
                    self._process_task(task_data)

            except redis.RedisError as e:
                self.logger.error(f"Redis error: {e}. Retrying in 5 seconds...")
                time.sleep(5)  # 重试前等待
            except Exception as e:
                self.logger.error(f"Unexpected error in queue listener: {e}")

    def _process_task(self, task_data):
        """解析并处理 Redis 任务"""
        try:
            task = json.loads(task_data)
            topic = task.get("send_topic")
            payload = task.get("send_data")
            
            if not topic or not payload:
                self.logger.warning(f"Invalid task format: {task_data[:100]}...")
                return
            
            # 转换 payload 为 JSON 字符串（如果是字典）
            if isinstance(payload, dict):
                payload = json.dumps(payload)
            
            # 发布消息（QoS 1 确保至少送达一次）
            result = self.client.publish(topic, payload, qos=1)
            if result[0] == 0:
                self.logger.info(f"Published to {topic}: {payload[:100]}...")

            elif result[0] == 4:
                self.logger.error("发布失败：MQTT 客户端未连接！")
                # 尝试重新连接并重发
                self.client.reconnect()
                self.client.publish(topic, payload, qos=1)  # 如果为 错误码为4时 可能是长时间没有活动client丢失，重新连接后可以重新发送
            else:
                self.logger.warning(f"Publish failed (status code {result[0]})")
                
        except json.JSONDecodeError as e:
            self.logger.error(f"JSON decode error: {task_data[:100]}... Error: {e}")
        except Exception as e:
            self.logger.error(f"Task processing failed: {e}")

    def stop(self):
        """优雅停止服务"""
        self.running = False
        if self.queue_thread and self.queue_thread.is_alive():
            self.queue_thread.join(5)  # 等待线程结束（最多 5 秒）
        
        self.client.loop_stop()
        self.client.disconnect()
        self.logger.info("MQTT Publisher stopped")


    

if __name__ == '__main__':

    if len(sys.argv) != 2:
        print("Usage: python mqtt_module.py <publisher|subscriber>")
        sys.exit(1)

    if sys.argv[1] == 'publisher':
        publisher = MqttPublisher()
        publisher.start()
    elif sys.argv[1] =='subscriber':
        subscriber = MqttSubscriber()
        subscriber.start()
    else:
        print("Invalid argument, please choose <publisher|subscriber>")
        sys.exit(1)
